import json
from typing import Dict, List, Optional, Any, Union
import logging
from datetime import datetime
import uuid
from exception.base_exceptions import DataProcessingException
from exception.exception_handler import global_exception_handler

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DataStorageHelper:
    """
    数据存储助手类
    用于将JSON数据存储到数据库和从数据库解析JSON数据
    """
    
    @staticmethod
    def create_session_data_for_storage(session_id: str, 
                                      wave_data: Dict[str, Any],
                                      delta_data: Dict[str, Any],
                                      theta_data: Dict[str, Any],
                                      alpha_data: Dict[str, Any],
                                      beta_data: Dict[str, Any],
                                      gamma_data: Dict[str, Any],
                                      metadata: Dict[str, Any],
                                      total_seconds: int) -> Dict[str, Any]:
        """
        创建用于存储的会话数据
        
        Args:
            session_id: 会话ID
            wave_data: 原始脑电数据
            delta_data: Delta波数据
            theta_data: Theta波数据
            alpha_data: Alpha波数据
            beta_data: Beta波数据
            gamma_data: Gamma波数据
            metadata: 元数据
            total_seconds: 总时长（秒）
            
        Returns:
            Dict: 用于存储的完整数据字典
        """
        try:
            # 生成唯一ID
            record_id = str(uuid.uuid4())
            
            # 构建存储数据
            storage_data = {
                'id': record_id,
                'session_id': session_id,
                'total_seconds': total_seconds,
                'wave': json.dumps(wave_data, ensure_ascii=False),
                'delta_wave': json.dumps(delta_data, ensure_ascii=False),
                'theta_wave': json.dumps(theta_data, ensure_ascii=False),
                'alpha_wave': json.dumps(alpha_data, ensure_ascii=False),
                'beta_wave': json.dumps(beta_data, ensure_ascii=False),
                'gamma_wave': json.dumps(gamma_data, ensure_ascii=False),
                'metadata': json.dumps(metadata, ensure_ascii=False)
            }
            
            logger.info(f"创建存储数据成功，会话ID: {session_id}")
            return storage_data
            
        except Exception as e:
            storage_error = DataProcessingException(
                f"创建存储数据失败，会话ID: {session_id}, 错误: {e}"
            )
            global_exception_handler.handle_exception(storage_error)
            raise storage_error
    
    @staticmethod
    def parse_stored_data_to_json(stored_record: Dict[str, Any]) -> Dict[str, Any]:
        """
        将存储的记录解析为JSON数据
        
        Args:
            stored_record: 从数据库获取的记录
            
        Returns:
            Dict: 解析后的JSON数据
        """
        try:
            parsed_data = {
                'id': stored_record.get('id'),
                'session_id': stored_record.get('session_id'),
                'created_at': stored_record.get('created_at'),
                'total_seconds': stored_record.get('total_seconds'),
                'wave': json.loads(stored_record.get('wave', '{}')),
                'delta_wave': json.loads(stored_record.get('delta_wave', '{}')),
                'theta_wave': json.loads(stored_record.get('theta_wave', '{}')),
                'alpha_wave': json.loads(stored_record.get('alpha_wave', '{}')),
                'beta_wave': json.loads(stored_record.get('beta_wave', '{}')),
                'gamma_wave': json.loads(stored_record.get('gamma_wave', '{}')),
                'metadata': json.loads(stored_record.get('metadata', '{}'))
            }
            
            logger.info(f"解析存储数据成功，会话ID: {parsed_data.get('session_id')}")
            return parsed_data
            
        except Exception as e:
            parse_error = DataProcessingException(
                f"解析存储数据失败，记录ID: {stored_record.get('id', 'unknown')}, 错误: {e}"
            )
            global_exception_handler.handle_exception(parse_error)
            raise parse_error
    
    @staticmethod
    def create_simple_session_data(session_id: str, 
                                 wave_frames: List[List[int]],
                                 delta_frames: List[List[int]],
                                 theta_frames: List[List[int]],
                                 alpha_frames: List[List[int]],
                                 beta_frames: List[List[int]],
                                 gamma_frames: List[List[int]],
                                 total_seconds: int,
                                 device_info: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        创建简单的会话数据（使用标准格式）
        
        Args:
            session_id: 会话ID
            wave_frames: 原始数据帧 [[index, value], ...]
            delta_frames: Delta波数据帧
            theta_frames: Theta波数据帧
            alpha_frames: Alpha波数据帧
            beta_frames: Beta波数据帧
            gamma_frames: Gamma波数据帧
            total_seconds: 总时长（秒）
            device_info: 设备信息
            
        Returns:
            Dict: 用于存储的完整数据字典
        """
        try:
            # 构建标准格式的数据
            wave_data = {
                "total_seconds": total_seconds,
                "frame_rate": 128,
                "wave": {"frames": wave_frames}
            }
            
            delta_data = {
                "total_seconds": total_seconds,
                "frame_rate": 1,
                "wave": {"frames": delta_frames}
            }
            
            theta_data = {
                "total_seconds": total_seconds,
                "frame_rate": 1,
                "wave": {"frames": theta_frames}
            }
            
            alpha_data = {
                "total_seconds": total_seconds,
                "frame_rate": 1,
                "wave": {"frames": alpha_frames}
            }
            
            beta_data = {
                "total_seconds": total_seconds,
                "frame_rate": 1,
                "wave": {"frames": beta_frames}
            }
            
            gamma_data = {
                "total_seconds": total_seconds,
                "frame_rate": 1,
                "wave": {"frames": gamma_frames}
            }
            
            # 构建元数据
            metadata = {
                "device_info": device_info or {
                    "device_name": "NeuroSky MindWave",
                    "device_type": "EEG",
                    "sampling_rate": 512,
                    "transmission_rate": 128
                },
                "signal_quality": "good",
                "session_notes": "",
                "created_at": datetime.now().isoformat()
            }
            
            # 调用标准方法创建存储数据
            return DataStorageHelper.create_session_data_for_storage(
                session_id=session_id,
                wave_data=wave_data,
                delta_data=delta_data,
                theta_data=theta_data,
                alpha_data=alpha_data,
                beta_data=beta_data,
                gamma_data=gamma_data,
                metadata=metadata,
                total_seconds=total_seconds
            )
            
        except Exception as e:
            session_error = DataProcessingException(
                f"创建简单会话数据失败，会话ID: {session_id}, 错误: {e}"
            )
            global_exception_handler.handle_exception(session_error)
            raise session_error
    
    @staticmethod
    def extract_wave_frames_from_json(parsed_data: Dict[str, Any]) -> Dict[str, List[List[int]]]:
        """
        从解析的数据中提取各频段的数据帧
        
        Args:
            parsed_data: 解析后的数据
            
        Returns:
            Dict: 各频段的数据帧
        """
        try:
            frames = {
                'wave': parsed_data.get('wave', {}).get('wave', {}).get('frames', []),
                'delta': parsed_data.get('delta_wave', {}).get('wave', {}).get('frames', []),
                'theta': parsed_data.get('theta_wave', {}).get('wave', {}).get('frames', []),
                'alpha': parsed_data.get('alpha_wave', {}).get('wave', {}).get('frames', []),
                'beta': parsed_data.get('beta_wave', {}).get('wave', {}).get('frames', []),
                'gamma': parsed_data.get('gamma_wave', {}).get('wave', {}).get('frames', [])
            }
            
            return frames
            
        except Exception as e:
            extract_error = DataProcessingException(
                f"提取数据帧失败，会话ID: {parsed_data.get('session_id', 'unknown')}, 错误: {e}"
            )
            global_exception_handler.handle_exception(extract_error)
            return {band: [] for band in ['wave', 'delta', 'theta', 'alpha', 'beta', 'gamma']}
    
    @staticmethod
    def validate_storage_data(storage_data: Dict[str, Any]) -> Dict[str, bool]:
        """
        验证存储数据的完整性
        
        Args:
            storage_data: 存储数据
            
        Returns:
            Dict: 验证结果
        """
        try:
            required_fields = ['session_id', 'total_seconds', 'wave', 'delta_wave', 
                             'theta_wave', 'alpha_wave', 'beta_wave', 'gamma_wave', 'metadata']
            
            validation_results = {}
            
            # 检查必需字段
            for field in required_fields:
                validation_results[f'has_{field}'] = field in storage_data and storage_data[field] is not None
            
            # 检查JSON格式
            json_fields = ['wave', 'delta_wave', 'theta_wave', 'alpha_wave', 'beta_wave', 'gamma_wave', 'metadata']
            for field in json_fields:
                try:
                    if field in storage_data and storage_data[field]:
                        json.loads(storage_data[field])
                        validation_results[f'{field}_json_valid'] = True
                    else:
                        validation_results[f'{field}_json_valid'] = False
                except:
                    validation_results[f'{field}_json_valid'] = False
            
            # 检查数据类型
            validation_results['total_seconds_is_int'] = isinstance(storage_data.get('total_seconds'), int)
            validation_results['session_id_is_string'] = isinstance(storage_data.get('session_id'), str)
            
            return validation_results
            
        except Exception as e:
            validation_error = DataProcessingException(
                f"验证存储数据失败，会话ID: {storage_data.get('session_id', 'unknown')}, 错误: {e}"
            )
            global_exception_handler.handle_exception(validation_error)
            return {'validation_failed': True}


# 便捷函数
def prepare_data_for_storage(session_id: str, 
                           wave_data: Dict[str, Any],
                           delta_data: Dict[str, Any],
                           theta_data: Dict[str, Any],
                           alpha_data: Dict[str, Any],
                           beta_data: Dict[str, Any],
                           gamma_data: Dict[str, Any],
                           metadata: Dict[str, Any],
                           total_seconds: int) -> Dict[str, Any]:
    """
    准备数据用于存储（便捷函数）
    
    Args:
        session_id: 会话ID
        wave_data: 原始脑电数据
        delta_data: Delta波数据
        theta_data: Theta波数据
        alpha_data: Alpha波数据
        beta_data: Beta波数据
        gamma_data: Gamma波数据
        metadata: 元数据
        total_seconds: 总时长（秒）
        
    Returns:
        Dict: 用于存储的完整数据字典
    """
    return DataStorageHelper.create_session_data_for_storage(
        session_id, wave_data, delta_data, theta_data, alpha_data, beta_data, gamma_data, metadata, total_seconds
    )


def parse_database_record(record: Dict[str, Any]) -> Dict[str, Any]:
    """
    解析数据库记录（便捷函数）
    
    Args:
        record: 数据库记录
        
    Returns:
        Dict: 解析后的JSON数据
    """
    return DataStorageHelper.parse_stored_data_to_json(record)


# 使用示例
if __name__ == "__main__":
    # 模拟从service层获取的数据
    wave_frames = [[i+1, 1000 + i] for i in range(100)]
    delta_frames = [[i+1, 200 + i] for i in range(10)]
    theta_frames = [[i+1, 300 + i] for i in range(10)]
    alpha_frames = [[i+1, 400 + i] for i in range(10)]
    beta_frames = [[i+1, 500 + i] for i in range(10)]
    gamma_frames = [[i+1, 600 + i] for i in range(10)]
    
    # 创建存储数据
    storage_data = DataStorageHelper.create_simple_session_data(
        session_id="test_session_001",
        wave_frames=wave_frames,
        delta_frames=delta_frames,
        theta_frames=theta_frames,
        alpha_frames=alpha_frames,
        beta_frames=beta_frames,
        gamma_frames=gamma_frames,
        total_seconds=100,
        device_info={"device_name": "Test Device"}
    )
    
    logger.info("存储数据创建成功:")
    logger.info(f"  ID: {storage_data['id']}")
    logger.info(f"  会话ID: {storage_data['session_id']}")
    logger.info(f"  总时长: {storage_data['total_seconds']}秒")
    
    # 验证数据
    validation = DataStorageHelper.validate_storage_data(storage_data)
    logger.info(f"数据验证结果: {validation}")
    
    # 模拟从数据库读取数据
    parsed_data = DataStorageHelper.parse_stored_data_to_json(storage_data)
    logger.info(f"数据解析成功，会话ID: {parsed_data['session_id']}")
    
    # 提取数据帧
    frames = DataStorageHelper.extract_wave_frames_from_json(parsed_data)
    logger.info(f"提取的数据帧:")
    for band, frame_data in frames.items():
        logger.info(f"  {band}: {len(frame_data)} 个数据点")